'''
本模块用于将pdf|doc|ppt|docx|pptx|html|txt|rtf|odt|xls|xlsx 转换为markdown，暂时不支持图片，因为没有多模态模型
HTML文档处理工具：pyppeteer
RTF文档处理工具：python-rtf
ODT文档处理工具：odfpy
XLS|XLSX文档处理工具：pandas
PDF|DOC|PPT|DOCX|PPTX 文档处理工具：MinerU(ocr)
封装为api接口，供其他模块调用，因为异步函数不能在主线程中使用
后续可以添加对网页url的解析
'''

import os
import shutil
import subprocess
import asyncio
import nest_asyncio
import pandas as pd
from flask import Flask, request, jsonify
from pyppeteer import launch
from striprtf.striprtf import rtf_to_text  # 使用 striprtf 替换 pyth
from odf.opendocument import load as load_odt
from odf.text import P

# 初始化 Flask 应用
app = Flask(__name__)

# 允许在已有事件循环中嵌套运行新的事件循环
nest_asyncio.apply()

# 支持的文件类型
SUPPORTED_TYPES = {
    '.pdf', '.doc', '.ppt', '.docx', '.pptx', '.html',
    '.txt', '.rtf', '.odt', '.xls', '.xlsx'
}

# 临时目录和输出目录（请根据实际环境修改路径）
TMP_DIR = './tmp_output'
OUTPUT_DIR = './db_files'

async def convert_html_to_pdf(input_path=None, input_url=None, output_dir=TMP_DIR):
    """使用 pyppeteer 将 HTML 转换为 PDF"""
    try:
        os.makedirs(output_dir, exist_ok=True)
        browser = await launch(
            handleSIGINT=False,
            handleSIGTERM=False,
            handleSIGHUP=False,
            headless=True,
            args=['--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage', '--disable-gpu']
        )
        page = await browser.newPage()
        if input_path:
            with open(input_path, 'r', encoding='utf-8') as f:
                html_content = f.read()
            await page.setContent(html_content)
            pdf_path = os.path.join(output_dir, os.path.basename(input_path).replace('.html', '.pdf'))
        elif input_url:
            await page.goto(input_url)
            pdf_path = os.path.join(output_dir, 'output.pdf')
        else:
            raise ValueError("Either input_path or input_url must be provided")
        await page.pdf({'path': pdf_path, 'format': 'A4', 'printBackground': True})
        await browser.close()
        return pdf_path
    except Exception as e:
        print(f"HTML 转 PDF 失败：{str(e)}")
        raise

async def convert_to_markdown(input_path, tmp_dir=TMP_DIR, output_dir=OUTPUT_DIR, method='ocr', lang=None, start_page=None, end_page=None):
    """转换文件到 Markdown 格式"""
    try:
        os.makedirs(tmp_dir, exist_ok=True)
        os.makedirs(output_dir, exist_ok=True)
        file_type = os.path.splitext(input_path)[-1].lower()

        if file_type not in SUPPORTED_TYPES:
            raise ValueError(f"不支持的文件类型：{file_type}")

        # 处理 HTML 文件：转换为 PDF，然后使用MinerU转换为markdown
        if file_type == '.html':
            pdf_path = await convert_html_to_pdf(input_path=input_path, output_dir=tmp_dir)
            if pdf_path and os.path.exists(pdf_path):
                input_path = pdf_path
                print(f"HTML 转 PDF 成功，存入 {pdf_path}")
                # 使用MinerU转换为markdown，递归调用
                return await convert_to_markdown(input_path, tmp_dir, output_dir, method, lang, start_page, end_page)
            else:
                raise ValueError("PDF conversion failed")

        # 处理 TXT 文件：直接读取并输出为 Markdown
        elif file_type == '.txt':
            with open(input_path, 'r', encoding='utf-8') as f:
                content = f.read()
            markdown_content = content
            file_name = os.path.basename(input_path).replace('.txt', '.md')
            with open(os.path.join(output_dir, file_name), 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            return True

        # 处理 RTF 文件：使用 striprtf 转换为纯文本
        elif file_type == '.rtf':
            with open(input_path, 'r', encoding='utf-8') as f:
                rtf_content = f.read()
            content = rtf_to_text(rtf_content)
            markdown_content = content
            file_name = os.path.basename(input_path).replace('.rtf', '.md')
            with open(os.path.join(output_dir, file_name), 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            return True

        # 处理 ODT 文件：使用 odfpy 提取文本
        elif file_type == '.odt':
            doc = load_odt(input_path)
            texts = [p.text for p in doc.getElementsByType(P) if p.text]
            content = '\n'.join(texts)
            markdown_content = content
            file_name = os.path.basename(input_path).replace('.odt', '.md')
            with open(os.path.join(output_dir, file_name), 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            return True

        # 处理 Excel 文件（.xls 和 .xlsx）
        elif file_type in ['.xls', '.xlsx']:
            df = pd.read_excel(input_path, sheet_name=None)
            markdown_content = ""
            for sheet_name, sheet_df in df.items():
                markdown_table = sheet_df.to_markdown(index=False)
                markdown_content += f"## {sheet_name}\n\n{markdown_table}\n\n"
            file_name = os.path.basename(input_path).replace(file_type, '.md')
            with open(os.path.join(output_dir, file_name), 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            return True

        # 其他文件类型（PDF、DOC、PPT 等）：使用 MinerU 转换
        else:
            command = ['magic-pdf', '-p', input_path, '-o', tmp_dir, '-m', method]
            if lang:
                command.extend(['-l', lang])
            if start_page is not None:
                command.extend(['-s', str(start_page)])
            if end_page is not None:
                command.extend(['-e', str(end_page)])
            subprocess.run(['conda', 'run', '-n', 'mineru'] + command, check=True)
            file_name = os.path.basename(input_path).split('.')[0]
            markdown_dir = os.path.join(tmp_dir, file_name, method)
            if os.path.exists(markdown_dir):
                for file in os.listdir(markdown_dir):
                    if file.endswith('.md'):
                        shutil.move(os.path.join(markdown_dir, file), os.path.join(output_dir, file))
                        print(f"Markdown 文件 {file} 移动至 {output_dir}")
            for file in os.listdir(tmp_dir):
                file_path = os.path.join(tmp_dir, file)
                if os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            return True

    except Exception as e:
        print(f"文件转换失败：{str(e)}")
        raise

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({'status': 'healthy', 'message': '服务正常运行'})

@app.route('/convert', methods=['POST'])
def convert_document():
    """
    文档转换接口
    请求参数 (JSON):
    {
        "input_path": "文件路径",
        "method": "ocr",  # 可选，默认为 ocr
        "lang": "语言",  # 可选
        "start_page": 1,  # 可选
        "end_page": 10   # 可选
    }
    """
    try:
        data = request.get_json()
        if not data or 'input_path' not in data:
            return jsonify({'status': 'error', 'message': '缺少必要参数：input_path'}), 400
        input_path = data['input_path']
        if not os.path.exists(input_path):
            return jsonify({'status': 'error', 'message': f'文件不存在：{input_path}'}), 404
        method = data.get('method', 'ocr')
        lang = data.get('lang')
        start_page = data.get('start_page')
        end_page = data.get('end_page')
        result = asyncio.run(convert_to_markdown(
            input_path=input_path,
            method=method,
            lang=lang,
            start_page=start_page,
            end_page=end_page
        ))
        return jsonify({'status': 'success', 'message': '文档转换成功', 'result': result})
    except Exception as e:
        return jsonify({'status': 'error', 'message': f'转换失败：{str(e)}'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)